

using namespace AscendC;


template <typename T>
class TokenShiftB1
{
public:
    __aicore__ inline TokenShiftB1() {}
    __aicore__ inline void Init(
        GM_ADDR x, 
        GM_ADDR rwkvag, 
        GM_ADDR h0, 
        GM_ADDR xr, 
        GM_ADDR xw, 
        GM_ADDR xk,
        GM_ADDR xv, 
        GM_ADDR xa, 
        GM_ADDR xg, 
        GM_ADDR ht, 
        GM_ADDR workspace,
        uint32_t aivNum,
        uint32_t B_length,
        uint32_t T_length,
        uint32_t C_length,
        uint32_t tiling_size,
        uint32_t core_rows,
        uint32_t core_rows_remain,
        uint32_t X_length_align,
        TPipe* pipeIn)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero");
        this->pipe = pipeIn;
        this->aivNum = aivNum;
        this->startPointer = core_rows * C_length * GetBlockIdx();
        this->core_rows = core_rows + (GetBlockNum() == GetBlockIdx() + 1 ? core_rows_remain : 0);
        this->B_length = B_length;
        this->T_length = T_length;
        this->C_length = C_length;
        this->tiling_size = tiling_size;
        this->ALIGN_NUM = 32 / sizeof(T);
        uint32_t core_length = this->core_rows * C_length + ALIGN_NUM;
        uint32_t h0startPointer = core_rows * C_length * GetBlockIdx();

        if(GetBlockIdx() == 0) {
            x_global.SetGlobalBuffer((__gm__ T *)x, core_length);
        } else {
            x_global.SetGlobalBuffer((__gm__ T *)x + this->startPointer - C_length, core_length + C_length);
        }
        rwkvag_global.SetGlobalBuffer((__gm__ T *)rwkvag, 6 * C_length + ALIGN_NUM);
        h0_global.SetGlobalBuffer((__gm__ T *)h0 , C_length + ALIGN_NUM);
        
        xr_global.SetGlobalBuffer((__gm__ T *)xr + this->startPointer, core_length);   
        xw_global.SetGlobalBuffer((__gm__ T *)xw + this->startPointer, core_length);   
        xk_global.SetGlobalBuffer((__gm__ T *)xk + this->startPointer, core_length);   
        xv_global.SetGlobalBuffer((__gm__ T *)xv + this->startPointer, core_length);   
        xa_global.SetGlobalBuffer((__gm__ T *)xa + this->startPointer, core_length);   
        xg_global.SetGlobalBuffer((__gm__ T *)xg + this->startPointer, core_length);        
        ht_global.SetGlobalBuffer((__gm__ T *)ht, C_length + ALIGN_NUM); 
        pipe->InitBuffer(inQueueXBind, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueRwkvag, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueH0, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueXNextLine, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXR, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXW, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXK, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXV, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXA, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXG, BUFFER_NUM, this->tiling_size * sizeof(T));

        pipe->InitBuffer(QueueTmp1, this->tiling_size * sizeof(T));
        g_xLocal = QueueTmp1.Get<T>();
    }

    __aicore__ inline void Process() {
        int32_t loopcount = this->C_length / this->tiling_size;
        uint32_t length_remain = this->C_length % this->tiling_size;
        uint32_t length_remain_align = (length_remain + this->ALIGN_NUM - 1) / this->ALIGN_NUM * this->ALIGN_NUM;


        for(int32_t p = 0; p < loopcount; p++) {
            if(GetBlockIdx() == 0) {
                CopyInH0(p * this->tiling_size, this->tiling_size);
                for(int32_t k = 0; k < this->core_rows; k++) {
                    CopyInXNextLine(k * this->C_length + p * this->tiling_size, this->tiling_size);
                    Compute(p * this->tiling_size, this->tiling_size);
                    CopyOutX(k * this->C_length + p * this->tiling_size, this->tiling_size);
                }
            } else {
                CopyInXPrevLine(p * this->tiling_size, this->tiling_size);
                for(int32_t k = 0; k < this->core_rows; k++) {
                    CopyInXNextLine(this->C_length + k * this->C_length + p * this->tiling_size, this->tiling_size);
                    Compute(p * this->tiling_size, this->tiling_size);
                    CopyOutX(k * this->C_length + p * this->tiling_size, this->tiling_size);
                }
            }
            if(this->aivNum == GetBlockIdx() + 1) {
                CopyOutHT(p * this->tiling_size, this->tiling_size);
            }
        }
        if(length_remain > 0) {
            if(GetBlockIdx() == 0) {
                CopyInH0(loopcount * this->tiling_size, length_remain_align);
                for(int32_t k = 0; k < this->core_rows; k++) {
                    CopyInXNextLine(k * this->C_length + loopcount * this->tiling_size, length_remain_align);
                    Compute(loopcount * this->tiling_size, length_remain_align);
                    if(length_remain == length_remain_align) {
                        CopyOutX(k * this->C_length + loopcount * this->tiling_size, length_remain);
                    } else {
                        CopyOutXUnalign(k * this->C_length + loopcount * this->tiling_size, length_remain);
                    }    
                }
            } else {
                CopyInXPrevLine(loopcount * this->tiling_size, length_remain_align);
                for(int32_t k = 0; k < this->core_rows; k++) {
                    CopyInXNextLine(this->C_length + k * this->C_length + loopcount * this->tiling_size, length_remain_align);
                    Compute(loopcount * this->tiling_size, length_remain_align);
                    if(length_remain == length_remain_align) {
                        CopyOutX(k * this->C_length + loopcount * this->tiling_size, length_remain);
                    } else {
                        CopyOutXUnalign(k * this->C_length + loopcount * this->tiling_size, length_remain);
                    }    
                }
            }
            if(this->aivNum == GetBlockIdx() + 1) {
                if(length_remain == length_remain_align) {
                    CopyOutHT(loopcount * this->tiling_size, length_remain);
                } else {
                    CopyOutHTUnalign(loopcount * this->tiling_size, length_remain);
                }    
            }
        }
               
    }
    
    __aicore__ inline void Compute(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xNextLineLocal = inQueueXNextLine.DeQue<T>();
        // LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        LocalTensor<T> XRLocal = outQueueXR.AllocTensor<T>();
        LocalTensor<T> XWLocal = outQueueXW.AllocTensor<T>();
        LocalTensor<T> XKLocal = outQueueXK.AllocTensor<T>();
        LocalTensor<T> XVLocal = outQueueXV.AllocTensor<T>();
        LocalTensor<T> XALocal = outQueueXA.AllocTensor<T>();
        LocalTensor<T> XGLocal = outQueueXG.AllocTensor<T>();

        Sub(g_xLocal, g_xLocal, xNextLineLocal, length);

        CopyInRline(startpos, length);
        LocalTensor<T> rwkvagLocal0 = inQueueRwkvag.DeQue<T>();
        Mul(XRLocal, g_xLocal, rwkvagLocal0, length);
        Add(XRLocal, XRLocal, xNextLineLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal0);
        outQueueXR.EnQue(XRLocal);  

        CopyInRline(startpos + this->C_length, length);
        LocalTensor<T> rwkvagLocal1 = inQueueRwkvag.DeQue<T>();
        Mul(XWLocal, g_xLocal, rwkvagLocal1, length);
        Add(XWLocal, XWLocal, xNextLineLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal1);
        outQueueXW.EnQue(XWLocal);  
        
        CopyInRline(startpos + this->C_length * 2, length);
        LocalTensor<T> rwkvagLocal2 = inQueueRwkvag.DeQue<T>();
        Mul(XKLocal, g_xLocal, rwkvagLocal2, length);
        Add(XKLocal, XKLocal, xNextLineLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal2);
        outQueueXK.EnQue(XKLocal);  

        CopyInRline(startpos + this->C_length * 3, length);
        LocalTensor<T> rwkvagLocal3 = inQueueRwkvag.DeQue<T>();
        Mul(XVLocal, g_xLocal, rwkvagLocal3, length);
        Add(XVLocal, XVLocal, xNextLineLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal3);
        outQueueXV.EnQue(XVLocal); 

        CopyInRline(startpos + this->C_length * 4, length);
        LocalTensor<T> rwkvagLocal4 = inQueueRwkvag.DeQue<T>();
        Mul(XALocal, g_xLocal, rwkvagLocal4, length);
        Add(XALocal, XALocal, xNextLineLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal4);
        outQueueXA.EnQue(XALocal); 

        CopyInRline(startpos + this->C_length * 5, length);
        LocalTensor<T> rwkvagLocal5 = inQueueRwkvag.DeQue<T>();
        Mul(XGLocal, g_xLocal, rwkvagLocal5, length);
        Add(XGLocal, XGLocal, xNextLineLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal5);
        outQueueXG.EnQue(XGLocal);      
        
        Adds(g_xLocal, xNextLineLocal, T(0), length);

        // inQueueX.EnQue(xLocal);
        inQueueXNextLine.FreeTensor(xNextLineLocal);
    }    

    __aicore__ inline void CopyInH0(uint32_t startpos, uint32_t length) {
        // LocalTensor<T> h0Local = inQueueH0.AllocTensor<T>();
        DataCopy(g_xLocal, h0_global[startpos], length);   
        PipeBarrier<PIPE_ALL>(); 
        // inQueueH0.EnQue(h0Local);         
    }       
    __aicore__ inline void CopyInXPrevLine(uint32_t startpos, uint32_t length) {
        DataCopy(g_xLocal, x_global[startpos], length);   
        PipeBarrier<PIPE_ALL>();      
    }       

    __aicore__ inline void CopyInXNextLine(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xLocal = inQueueXNextLine.AllocTensor<T>();
        DataCopy(xLocal, x_global[startpos], length);   
        inQueueXNextLine.EnQue(xLocal);  
    }

    __aicore__ inline void CopyInRline(uint32_t startpos, uint32_t length) {
        LocalTensor<T> rwkvagLocal = inQueueRwkvag.AllocTensor<T>();
        DataCopy(rwkvagLocal, rwkvag_global[startpos], length);   
        inQueueRwkvag.EnQue(rwkvagLocal);  
    }

    __aicore__ inline void CopyOutHT(uint32_t startpos, uint32_t length) {
        // LocalTensor<T> xLocal = inQueueXBind.DeQue<T>();
        DataCopy(ht_global[startpos], g_xLocal, length);  
        PipeBarrier<PIPE_ALL>();  
    }

    __aicore__ inline void CopyOutHTUnalign(uint32_t startpos, uint32_t length) {
        uint32_t process_length = length / this->ALIGN_NUM * this->ALIGN_NUM;

        DataCopy(ht_global[startpos], g_xLocal, process_length); 
        for(int32_t m = process_length; m < length; m++) {
            ht_global.SetValue(m + startpos, g_xLocal.GetValue(m));
        }    
        PipeBarrier<PIPE_ALL>();

    }

    __aicore__ inline void CopyOutX(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        DataCopy(xr_global[startpos], xrLocal, length); 
        outQueueXR.FreeTensor(xrLocal);

        LocalTensor<T> xwLocal = outQueueXW.DeQue<T>();
        DataCopy(xw_global[startpos], xwLocal, length); 
        outQueueXW.FreeTensor(xwLocal);

        LocalTensor<T> xkLocal = outQueueXK.DeQue<T>();
        DataCopy(xk_global[startpos], xkLocal, length); 
        outQueueXK.FreeTensor(xkLocal);

        LocalTensor<T> xvLocal = outQueueXV.DeQue<T>();
        DataCopy(xv_global[startpos], xvLocal, length); 
        outQueueXV.FreeTensor(xvLocal);

        LocalTensor<T> xaLocal = outQueueXA.DeQue<T>();
        DataCopy(xa_global[startpos], xaLocal, length); 
        outQueueXA.FreeTensor(xaLocal);

        LocalTensor<T> xgLocal = outQueueXG.DeQue<T>();
        DataCopy(xg_global[startpos], xgLocal, length); 
        outQueueXG.FreeTensor(xgLocal);
        PipeBarrier<PIPE_ALL>();
    }
    __aicore__ inline void CopyOutXUnalign(uint32_t startpos, uint32_t length) {
        uint32_t process_length = length / this->ALIGN_NUM * this->ALIGN_NUM;

        LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        DataCopy(xr_global[startpos], xrLocal, process_length); 
        for(int32_t m = process_length; m < length; m++) {
            xr_global.SetValue(m + startpos, xrLocal.GetValue(m));
        }    
        PipeBarrier<PIPE_ALL>();
        outQueueXR.FreeTensor(xrLocal);

        LocalTensor<T> xwLocal = outQueueXW.DeQue<T>();
        DataCopy(xw_global[startpos], xwLocal, process_length); 
        for(int32_t m = process_length; m < length; m++) {
            xw_global.SetValue(m + startpos, xwLocal.GetValue(m));
        }     
        PipeBarrier<PIPE_ALL>();
        outQueueXW.FreeTensor(xwLocal);

        LocalTensor<T> xkLocal = outQueueXK.DeQue<T>();
        DataCopy(xk_global[startpos], xkLocal, process_length); 
        for(int32_t m = process_length; m < length; m++) {
            xk_global.SetValue(m + startpos, xkLocal.GetValue(m));
        }    
        PipeBarrier<PIPE_ALL>();
        outQueueXK.FreeTensor(xkLocal);

        LocalTensor<T> xvLocal = outQueueXV.DeQue<T>();
        DataCopy(xv_global[startpos], xvLocal, process_length); 
        for(int32_t m = process_length; m < length; m++) {
            xv_global.SetValue(m + startpos, xvLocal.GetValue(m));
        }      
        PipeBarrier<PIPE_ALL>();
        outQueueXV.FreeTensor(xvLocal);

        LocalTensor<T> xaLocal = outQueueXA.DeQue<T>();
        DataCopy(xa_global[startpos], xaLocal, process_length); 
        for(int32_t m = process_length; m < length; m++) {
            xa_global.SetValue(m + startpos, xaLocal.GetValue(m));
        }    
        PipeBarrier<PIPE_ALL>();
        outQueueXA.FreeTensor(xaLocal);

        LocalTensor<T> xgLocal = outQueueXG.DeQue<T>();
        DataCopy(xg_global[startpos], xgLocal, process_length); 
        for(int32_t m = process_length; m < length; m++) {
            xg_global.SetValue(m + startpos, xgLocal.GetValue(m));
        }     
        PipeBarrier<PIPE_ALL>();
        outQueueXG.FreeTensor(xgLocal);
    }
private:
    GlobalTensor<T> x_global;
    GlobalTensor<T> rwkvag_global;
    GlobalTensor<T> h0_global;
    GlobalTensor<T> xr_global, xw_global, xk_global, xv_global, xa_global, xg_global, ht_global; 

    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQueueXBind;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueRwkvag, inQueueH0;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueXNextLine;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueXR, outQueueXW, outQueueXK, outQueueXV, outQueueXA, outQueueXG;

    TBuf<QuePosition::VECCALC> QueueTmp1;
    uint32_t B_length;
    uint32_t T_length;
    uint32_t C_length;
    uint32_t tiling_size;
    uint32_t core_rows;
    uint32_t core_rows_remain;
    uint32_t startPointer;
    uint32_t X_length_align;
    uint32_t ALIGN_NUM;
    uint32_t aivNum;

    LocalTensor<T> g_xLocal;
    AscendC::TPipe* pipe;
};